# pip3 install transformers fastapi uvicorn
# python3 app.py

import transformers
from fastapi import FastAPI
from pydantic import BaseModel


# 定义请求体模型
class TextInput(BaseModel):
    text: str


# 初始化 tokenizer
chat_tokenizer_dir = "./"
tokenizer = transformers.AutoTokenizer.from_pretrained(chat_tokenizer_dir, trust_remote_code=True)

# 初始化 FastAPI 应用
app = FastAPI()


@app.post("/tokenize")
async def tokenize_text(input: TextInput):
    """
        计算输入文本的 token 长度

        参数:
        - text: 需要计算 token 长度的文本内容

        返回:
        - token_count: token 数量
        - token_ids: token 的 ID 列表
        """
    # 使用 tokenizer 编码文本
    result = tokenizer.encode(input.text)

    # 返回 token 数量和 token ID 列表
    return {
        "token_count": len(result),
        "token_ids": result
    }


if __name__ == "__main__":
    import uvicorn

    uvicorn.run(app, host="0.0.0.0", port=8000)
